#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
sys.path.insert(0, '.')
import os


file_faceid_merge_list_path = r'/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/file_faceid_merge_list.txt'
file_list_path = r'/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/BoxFeatureMerge.norm.txt'
out_file_list_path = r'/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/FaceIDBoxFeatureMerge.norm.txt'


def main():
    face_ids = set()
    merge_video_faceids = {}
    max_faceid = 0
    with open(file_faceid_merge_list_path, 'r') as file:
        for line in file.readlines():
            # gd/zs/sqqjd/yxxhjt/150100414a54443452067fa6d4be6600/video/20190924_213328/00062_000.jpg 564
            lines = line.strip().split()
            file_name = lines[0]
            faceid = int(lines[1])
            video_name = os.path.dirname(file_name)
            # if merge_video_faceids.get(video_name, None) is not None:
            #     print(video_name)
            merge_video_faceids[video_name] = faceid
            face_ids.add(faceid)
            if max_faceid < faceid:
                max_faceid = faceid
    print('Merge FaceID Count:', len(face_ids))

    face_id = max_faceid
    video_faceids = {}
    faceid_image_counts = {}
    file_names = []
    with open(file_list_path, 'r') as file:
        for line in file.readlines():
            # gd/zs/dqjd/xgly/150100414a5444345203bcd04294b500/video/20190912_154308/00010_000.jpg 154,69,290,245,1,1024
            file_name = line.strip().split()[0]
            video_name = os.path.dirname(file_name)
            cur_face_id = video_faceids.get(video_name, None)
            if cur_face_id is None:
                cur_face_id = merge_video_faceids.get(video_name, None)
                if cur_face_id is None:
                    face_id += 1
                    cur_face_id = face_id
                video_faceids[video_name] = cur_face_id
            if faceid_image_counts.get(cur_face_id, None) is None:
                faceid_image_counts[cur_face_id] = 1
            else:
                faceid_image_counts[cur_face_id] += 1
            file_names.append(file_name)

    face_ids = set()
    with open(out_file_list_path, 'w') as file:
        for file_name in file_names:
            video_name = os.path.dirname(file_name)
            face_id = video_faceids[video_name]
            face_ids.add(face_id)
            image_count = faceid_image_counts[face_id]
            out_info = '{} {} {} {}\n'.format(face_id, face_id, file_name, image_count)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_list_path))
    print('FaceID Count:', len(face_ids))
    print('Finish!')


if __name__ == '__main__':
    main()
